import matplotlib.pyplot as plt
import torch
import numpy as np
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# 加载iris数据集
data, target = load_iris(return_X_y=True)

# 将数据转换为二分类问题
# 选择类别0和1的数据，排除类别2
data = data[target != 2]
target = target[target != 2]

# 数据标准化
scaler = StandardScaler()
data_scaled = scaler.fit_transform(data)

# 转换为PyTorch张量
x = torch.Tensor(data_scaled)
y = torch.Tensor(target).reshape(-1, 1)

# 划分训练集和测试集
train_x, test_x, train_y, test_y = train_test_split(x, y)

# 定义模型
# 为了匹配iris数据集的特征数量（4个特征），我们需要调整模型的输入特征数量
model = torch.nn.Sequential(
    torch.nn.Linear(in_features=4, out_features=1),  # 修改输入特征数量以匹配iris数据集
    torch.nn.Sigmoid()
)

# 损失函数和优化器
loss_fn = torch.nn.BCELoss()
op = torch.optim.RMSprop(params=model.parameters())

# 训练模型
loss_list = []
for i in range(2000):  # 可能需要调整迭代次数
    op.zero_grad()
    h = model(train_x)
    loss = loss_fn(h, train_y)
    loss_list.append(loss.item())
    loss.backward()
    op.step()
    if i % 100 == 0:
        print(f'Loss at iteration {i}: {loss.item()}')

# 测试模型
predict = model(test_x)
y_predict = (predict > 0.5).int()
acc = (y_predict == test_y).float().mean().item()
print(f'Accuracy: {acc}')
